import random
import copy

class DomainInterface:
    """A base class for an arbitrary domain"""
    def dimensionality(self):
        return 0

    def bounds(self):
        return None
    
    def contains(self,x):
        raise NotImplementedError()
    
    def sample(self):
        raise NotImplementedError()

class EmptyDomain(DomainInterface):
    def contains(self,x):
        return False
    
    def sample(self):
        return None

class DiscreteDomain(DomainInterface):
    """A domain consisting of a discrete list"""
    def __init__(self,items):
        self.items = items

    def __str__(self):
        return '{'+",".join([str(i) for i in self.items])+'}'

    def __iter__(self):
        return self.items.__iter__()

    def __len__(self):
        return len(self.items)

    def bounds(self):
        return (min(self.items),max(self.items))

    def contains(self,x):
        return x in self.items
    
    def sample(self):
        return random.choice(self.items)

class DiscreteRangeDomain(DomainInterface):
    """A domain consisting of a discrete integer range [a,b]"""
    def __init__(self,a,b):
        self.low = a
        self.high = b

    def __str__(self):
        return str(self.low)+':'+str(self.high)

    def __iter__(self):
        return xrange(self.low,self.high+1).__iter__()

    def __len__(self):
        return self.high+1-self.low

    def bounds(self):
        return (self.low,self.high)

    def contains(self,x):
        return isinstance(x,int) and x >= self.low and x <= self.high
    
    def sample(self):
        return random.randint(self.low,self.high)

class ContinuousDomain(DomainInterface):
    """A domain consisting of a continuous range [a,b]"""
    def __init__(self,*args):
        """Can take either a single pair argument or two arguments as input"""
        if len(args)==2:
            self.low = args[0]
            self.high = args[1]
        elif len(args)==1:
            assert len(args[0])==2,"Input must be a pair"
            self.low,self.high = args[0]
        else:
            raise ValueError("Input must be either a pair or two arguments")

    def __str__(self):
        return '('+str(self.low)+','+str(self.high)+')'

    def dimensionality(self):
        return 1

    def bounds(self):
        return (self.low,self.high)

    def contains(self,x):
        return x >= self.low and x <= self.high
    
    def sample(self):
        return random.uniform(self.low,self.high)

class CompositeDomain(DomainInterface):
    """A domain with several elements, given in an input list.
    For example, a multidimensional domain can be defined as
    [ContinuousDomain1,ContinuousDomain2,...,].
    """
    def __init__(self,args):
        self.elements = [Domain(a) for a in args]

    def __str__(self):
        return '['+','.join(str(e) for e in self.elements)+']'

    def isFlat(self):
        return not any(isinstance(e,CompositeDomain) for e in self.elements)

    def isMultivariate(self):
        return all(isinstance(e,ContinuousDomain) for e in self.elements)

    def dimensionality(self):
        return sum(e.dimensionality() for e in self.elements)

    def bounds(self):
        return [e.bounds() for e in self.elements]

    def contains(self,x):
        assert len(x) == len(self.elements),"Wrong dimensionality"
        for i in xrange(len(x)):
            if not self.elements[i].contains(x[i]):
                return False
        return True
    
    def sample(self):
        return [e.sample() for e in self.elements]

"""Can be initialized with either:
- a DomainInterface object
- a list of integer values (integer set)
- a list of domains (e.g., multidimensional continuous range)
- a 2-tuple (1-dimensional continuous range)
- an xrange object (integer range)
Arguments are checked for those types, in that order
"""
def Domain(arg):
    if isinstance(arg,DomainInterface):
        return arg
        #return copy.deepcopy(arg)
    elif isinstance(arg,list):
        if isinstance(arg[0],int):
            return DiscreteDomain(arg);
        else:
            return CompositeDomain(arg)
    elif isinstance(arg,tuple):
        assert(len(arg)==2),"Tuple argument must be a pair"
        return ContinuousDomain(arg[0],arg[1])
    elif isinstance(arg,xrange):
        if len(arg) > 1:
            assert(arg[1]-arg[0] == 1),"Domain xrange constructors must contain all integer in a range"
        return DiscreteRangeDomain(arg[0],arg[-1])
    else:
        raise ValueError("Unknown argument "+repr(arg)+" to Domain constructor")
